module CudaPrinter(printExpr, printFunction, printProgram) where

import Text.PrettyPrint.HughesPJ
import Types
import Data.Char(toUpper)
import Data.List(isSuffixOf)


printExpr :: Expr -> String
printExpr e = render (expr e)


expr :: Expr -> Doc
expr (Ci x)              = integer x
expr (C x)               = double x
expr (V x)               = text x
expr (S x)               = text x
expr (Op "plus" x y)     = prec 4 x <+> char '+' <+> prec 4 y
expr (Op "minus" x y)    = prec 4 x <+> char '-' <+> prec 5 y
expr (Op "mult" x y)     = prec 5 x <+> char '*' <+> prec 5 y
expr (Op "div" x y)      = prec 5 x <+> char '/' <+> prec 6 y
expr (Op f x y)          = text f <> (parens $ commas $ map expr $ [x, y])
expr (Func "lt" [x, y])  = expr x <+> char '<' <+> expr y
expr (Func "gt" [x, y])  = expr x <+> char '>' <+> expr y
expr (Func "eq" [x, y])  = prec 2 x <+> text "==" <+> prec 3 y
expr (Func "and" [x, y]) = prec 1 x <+> text "&&" <+> prec 1 y
expr (Func "or" [x, y])  = prec 0 x <+> text "||" <+> prec 0 y
expr (Func "not" [x])    = char '!' <> prec 10 x
expr (Func f as)         = text f <> (parens $ commas $ map expr $ as)
expr (Cast t e)          = text "make_" <> text (typeName t) <> (parens $ expr e)
expr (Dot e x)           = expr e <> char '.' <> text x
expr T                   = text "true"
expr F                   = text "false"


commas :: [Doc] -> Doc
commas es = hsep (punctuate (char ',')  es)

semis :: [Doc] -> Doc
semis es = hsep (punctuate semi es)


prec :: Int -> Expr -> Doc
prec n x = if n > p x then parens (expr x) else expr x where
    p (Func "not" _)   = 2
    p (Func "or" _)   = 0
    p (Func "and" _)  = 1

    p (Func "eq" _)   = 2
    p (Func "lt" _)   = 2
    p (Func "gt" _)   = 2

    p (Op "plus"_ _)  = 4
    p (Op "minus"_ _) = 4
    p (Op "mult"_ _)  = 5
    p (Op "div"_ _)   = 5
    
    p _ = 10


printProgram :: Program -> String
-- printProgram p = render (prefix $$ stmts p $$ suffix)
printProgram p = render (vcat $ map func p)


prefix = text "__device__ float3 evalRay(float3 S, float3 E)\n{"
suffix = char '}'

    
printFunction :: Function -> String
printFunction f = render (func f)


func :: Function -> Doc
func (Function f t as b) = special $$ text "GPU" <+> text (typeName t) <+> text f <>
    (parens $ commas $ map arg as) $$ stmtBlock b $+$ text ""
        where
            arg (x, t) = text (typeName t) <+> text x
            special = if f == "evalRgba" then text "#define EVAL_RGBA" else empty


stmtBlock (Group xs) = char '{' $+$ (nest 4 $ stmts xs) $+$ char '}'
stmtBlock s = char '{' $+$ (nest 4 $ stmts [s]) $+$ char '}'


stmts :: [Stmt] -> Doc
stmts [] = empty

-- parameter definitions

-- ugly hack; fix later to support true 1D textures
stmts (TypeDecl x Dataset : Assign y (Func "data1d_rgba" [S filename, S typename]) : rest) | x == y =
    text "DATA" <> parens (commas [text x, text filename, text typename, integer 2, integer 4]) $$ stmts rest
stmts (TypeDecl x Dataset : Assign y (Func "data1d" [S filename, S typename]) : rest) | x == y =
    text "DATA" <> parens (commas [text x, text filename, text typename, integer 2, integer 1]) $$ stmts rest

stmts (TypeDecl x Dataset : Assign y (Func "data2d_rgba" [S filename, S typename]) : rest) | x == y =
    text "DATA" <> parens (commas [text x, text filename, text typename, integer 2, integer 4]) $$ stmts rest
stmts (TypeDecl x Dataset : Assign y (Func "data2d" [S filename, S typename]) : rest) | x == y =
    text "DATA" <> parens (commas [text x, text filename, text typename, integer 2, integer 1]) $$ stmts rest

stmts (TypeDecl x Dataset : Assign y (Func "data3d" [S filename, S typename]) : rest) | x == y =
    text "DATA" <> parens (commas [text x, text filename, text typename, integer 3, integer 1]) $$ stmts rest

stmts (TypeDecl x Dataset : Assign y (Func "data4d" [S filename, S typename]) : rest) | x == y =
    text "DATA" <> parens (commas [text x, text filename, text typename, integer 4, integer 1]) $$ stmts rest

stmts (TypeDecl x t : Assign y (Func "float_param" args) : rest) | x == y =
    let
        f' = "PARAM_RANGE"
        args' = commas $ map expr args
    in text f' <> parens (text x <> char ',' <+> args') $$ stmts rest

-- stmts (TypeDecl x t : Assign y (Func "float3_param" args) : rest) | x == y =
--     let
--         f' = "FLOAT3"
--         args' = foldl1 (<+>) (map expr args)
--     in text f' <> parens (text x <> char ',' <+> args') $$ stmts rest

-- regular statements
stmts (Assign x e : rest) = (text x <+> equals <+> expr e <> char ';') $$ stmts rest
stmts (TypeDecl x t : Assign y e : rest) | x ==y  =
    (text (typeName t) <+> text x <+> equals <+> expr e <> char ';') $$ stmts rest
stmts (Return e _ : rest) = (text "return" <+> expr e <> char ';') $$ stmts rest
stmts (Branch c s (Group []) : rest) = (text "if" <+> parens (expr c) $$ stmtBlock s) $$ stmts rest
stmts (Branch c s1 s2 : rest) =
    (text "if" <+> parens (expr c) $$ stmtBlock s1 $$ text "else" $$ stmtBlock s2) $$ stmts rest
stmts (While c s : rest) = (text "while" <+> parens (expr c) $$ stmtBlock s) $$ stmts rest
stmts (For c (Assign x u) s : rest) =
    let
        upd = text x <+> equals <+> expr u
    in (text "for" <+> parens (semis [empty, expr c, upd]) $$ stmtBlock s) $$ stmts rest
stmts (FunDef f : rest) = func f $$ stmts rest
stmts (Break : rest) = text "break"  <> char ';' $$  stmts rest
stmts (Continue : rest) = text "continue"  <> char ';' $$  stmts rest


typeName :: Type -> String
typeName Float      = "float"
typeName Float2     = "float2"
typeName Float3     = "float3"
typeName Float4     = "float4"
typeName Boolean    = "bool"
typeName Int        = "int"
typeName Int2       = "int2"
typeName Int3       = "int3"
typeName Int4       = "int4"
typeName Dataset    = "data"
typeName NoType     = ""



